'''
Created on 2009-10-28

@author: Qiu wenfeng
'''
from jolly.imp.levelset.fspecial import *
from jolly.imp.levelset.lsecml05 import *
import vtk
import vtk.util.numpy_support as VN
import matplotlib.pyplot as plt
import sys
import scipy
import scipy.ndimage
import tempfile, os, shutil
import subprocess                 # For issuing commands to the OS.
if __name__ == '__main__':
   
    # read the image
    reader = vtk.vtkBMPReader()
    if len(sys.argv)>1:
        reader.SetFileName(sys.argv[1])
    else:
        reader.SetFileName('../data/twocells.bmp')
    reader.Update()
    
    il = vtk.vtkImageLuminance()
    il.SetInput(reader.GetOutput())
    il.Update()
    flip = vtk.vtkImageFlip()
    flip.SetFilteredAxis(1)
    flip.SetInput(il.GetOutput())
    flip.Update()
    
    # vtk to numpy
    imarray = VN.vtk_to_numpy(flip.GetOutput().GetPointData().GetScalars())
    xmin, xmax, ymin, ymax, zmin, zmax = flip.GetOutput().GetWholeExtent()
    imarray = imarray.reshape(ymax+1, xmax+1)
    
    
    G = fspecial('gaussian', N=15, Sigma=1.5)
    # smooth image by Gaussiin convolution
    imarray = scipy.ndimage.filters.convolve(imarray, G, output=float, mode='constant') # some diffent value from matlab
    y, x = scipy.gradient(imarray)
    f = y**2+x**2
   
    # edge indicator function.
    g = 1/(1+f)
    # the papramater in the definition of smoothed Dirac function
    epsilon=1.5
    # time step
    timestep=5
    # coefficient of the internal (penalizing) energy term P(\phi)
    # Note: the product timestep*mu must be less than 0.25 for stability!
    mu=0.2/timestep
    # coefficient of the weighted length term Lg(\phi)
    lmd=5
    # coefficient of the weighted area term Ag(\phi);
    # Note: Choose a positive(negative) alf if the initial contour is outside(inside) the object.
    alf=1.5
    
    # define initial level set function (LSF) as -c0, 0, c0 at points outside, on
    # the boundary, and inside of a region R, respectively.
    nrow, ncol = imarray.shape
    c0=4
    initialLSF=c0*scipy.ones((nrow, ncol), float)
    w=8
    # zero level set is on the boundary of R. 
    # Note: this can be commented out. The intial LSF does NOT necessarily need a zero level set.
    initialLSF[w:-w, w:-w]=0
    initialLSF[w+1:-w-1, w+1: -w-1]=-c0
    u=initialLSF
    plt.figure()
    
    
    plt.imshow(imarray)
    plt.gray()
    plt.contour(u,scipy.arange(-1, 0, 1), colors='r')
    plt.title('Initial contour')
#    plt.show()
    tmpdir = tempfile.mkdtemp()
    print tmpdir
    filenames = []
    filename = os.path.join(tmpdir,str('%03d' % 1) + '.png')
    filenames.append(filename+"\n")
    plt.savefig(filename, dpi=100)
    plt.clf()
    
    # start level set evolution
    for n in range(300):
        u = lsecml05(u, g ,lmd, mu, alf, epsilon, timestep, 1)
        if n%10==0:
            plt.close()
            plt.imshow(imarray)
            plt.gray()
            plt.contour(u,scipy.arange(-1, 0, 1), colors='r')
            plt.title('%d iterations'%(n+1))
#            plt.show()
            filename = os.path.join(tmpdir,str('%03d' % (n+2)) + '.png')
            filenames.append(filename+"\n")
            plt.savefig(filename, dpi=100)
            plt.clf()
    
    plt.imshow(imarray)
    plt.gray()
    plt.contour(u,scipy.arange(-1, 0, 1), colors='r')
    plt.title('%d iterations'%(n+1))
    plt.show()
    
    f = open('list.txt', 'w')
    f.writelines(filenames)
    f.close()
    
    command = ('..//../tools//mencoder',
           'mf://@list.txt',
           '-mf',
           'type=png:w=800:h=600:fps=25',
           '-ovc',
           'lavc',
           '-lavcopts',
           'vcodec=mpeg4:mbd=2:trell:autoaspect',
           '-oac',
           'copy',
           '-o',
           'output.avi')
    print "\n\nabout to execute:\n%s\n\n" % ' '.join(command)
    subprocess.check_call(command)
    
    shutil.rmtree(tmpdir)